#include "tunnel_server.h"


//read SOCKS identifier message
IdMessage * readIDMessage(int sock)
{  
  IdMessage * m = (IdMessage *)safeRead(sock,2);  
  m->methods = safeRead(sock,m->numMethods);    
  printf("handshake ver: %d, nmeth: %d, firstMethod: %d\n", m->version, m->numMethods, (m->methods)[0]);
  return m;
}

//checks if protocol version is 5, ak nie tak ho zrusi
void checkProtocol(int sock, int version)
{  
  if(version != 0x05) 
  {
	cerr << "Protocol version not supported" << endl;
	shutdown(sock, SHUT_RDWR);
	close(sock);
	exit(1);
  }
}

bool isSupported(int m) // skontrolujue ci podporujeme danu metodu
{
  if(m==0 /*|| m==1*/)
  {
    return true;
  }
  return false;
}

//select message and send reply
void writeIDReply(int sock, IdMessage * msg)
{ 
  //send reply
  char method = 0xff;
  for(int i = 0; i < msg->numMethods; i++)
  {
    if(isSupported(msg->methods[i]))
    {
      method = msg->methods[i];
      break;
    }
  }  
  uint8_t reply[] = {0x5,method};
  safeWrite(sock,reply,sizeof(reply));
}

//read SOCKS identifier message
SocksRequest * readRequest(int sock) //TODO : co ak pride nejaka haluz?
{ 
  SocksRequest * req = (SocksRequest * )safeRead(sock,4);  
  switch(req->atype)
  {
    case SAT_IPv4:
      req->size = 4;      
      break;
    case SAT_DOMAIN:
      req->size = *safeRead(sock,1);      
      break;
    case SAT_IPv6:
      req->size = 16;      
      break;
    default:
      req->size = 0;      
  }
  req->address = safeRead(sock,req->size);       
  req->port = htons(*((uint16_t *)safeRead(sock,2)));
  
  //Debug vypis:
  printf("Prisiel cmd: %x, adresa:",req->cmd);
  for(int i = 0;i<req->size;i++)
  {
    if(i>0) printf(".");
    printf("%d",req->address[i]);
  }
  printf(":%d\n",req->port);
  return req;
}

sockaddr_in getRecipient(int clientSock,SocksRequest * req)
{
   // recipient structure holds every necessary recipient info
  sockaddr_in recipient;
  recipient.sin_port = htons(req->port);
  switch(req->atype){
    case SAT_IPv4: // IPv4
      recipient.sin_family = AF_INET;
      //inet_pton(AF_INET, "127.0.0.1", req->address); // maybe for testing
      recipient.sin_addr.s_addr = *(unsigned long*)req->address;
      break;
    case SAT_DOMAIN: // Domain name
          struct addrinfo * adresa;
      recipient.sin_family = AF_INET;
      if(getaddrinfo((char *)req->address, NULL, NULL, &adresa) != 0)
        perror("getaddrinfo");
      printf("adresa: %s\n", inet_ntoa((*(struct sockaddr_in *)(adresa->ai_addr)).sin_addr)); // debug vypis
          recipient.sin_addr = (*(struct sockaddr_in *)(adresa->ai_addr)).sin_addr;
          freeaddrinfo(adresa);
      break;
    case SAT_IPv6: 
      recipient.sin_family = AF_INET6;
      recipient.sin_addr.s_addr = (unsigned long)req->address;
      break;
    default:
      cerr << "Unable to determine target address" << endl;
      shutdown(clientSock, SHUT_RDWR);
      close(clientSock);
      exit(1); // FIXME exit code
  }
  // so we have recipient family, address and port sucessfuly filled in
  return recipient;
}

void writeReplay(int sock, uint8_t repcode,uint8_t* addr = NULL , uint8_t addrsz = 0 , uint16_t port = 0)
{   
    //nastavenie reply message
    size_t len = 4+addrsz+sizeof(port);
    SocksReply * rpl = (SocksReply*) malloc(len);
    memset(rpl,0,len);
    
    rpl->version = 5;
    rpl->rep = repcode;
    rpl->rsv = 0;      
    if(addrsz==4)
      rpl->atype = SAT_IPv4;
    else
      rpl->atype = SAT_IPv6;
    //adresa
    uint8_t * address = (uint8_t*) rpl->address_data;
    memcpy(address,addr,addrsz);  
    //port
    address+=addrsz;
    memcpy(address,&port,sizeof(port));    
    //send reply
    safeWrite(sock,rpl,len);
}



void reqConnect(int clientSock,SocksRequest * req)
{
  
  sockaddr_in recipient = getRecipient(clientSock,req); 
  
  int outSock = socket(AF_INET, SOCK_STREAM, 0);  
  
  //connect
  if(connect(outSock,(sockaddr *)&recipient,sizeof(recipient))!=-1)
  {
    cout << "Connected" << endl;
    sockaddr sa;
    size_t sa_len = sizeof(sa);
    if(getsockname(outSock,(sockaddr *)&sa,(socklen_t *)&sa_len)<0)   
    {
        perror("getsockname");
        exit(1);
    }   
    
    if(sa.sa_family ==  AF_INET)
    {
      sockaddr_in * sain = (sockaddr_in *) &sa;
      //nastavenie replay adresy(BIND.ADDR)    
      uint8_t* addr = (uint8_t * )sain->sin_addr.s_addr;  
      
      //nastavenie replay portu(BIND.PORT)
      uint16_t port = sain->sin_port;
      
      //success reply
      writeReplay(clientSock,0,addr,4,port);
    }    
    else if(sa.sa_family ==  AF_INET6)
    {
      sockaddr_in6 * sain = (sockaddr_in6 *) &sa;
      //nastavenie replay adresy(BIND.ADDR)    
      uint8_t* addr = (uint8_t * )sain->sin6_addr.s6_addr;  
      
      //nastavenie replay portu(BIND.PORT)
      uint16_t port = sain->sin6_port;
      
      //success reply
      writeReplay(clientSock,0,addr,16,port);
    }
    
    exchangeData(clientSock,outSock);
    
  }
  else
  {    
    //send failure reply    
    writeReplay(clientSock,1); //FIXME: vybrat spravny chybovy kod, zatial budem posielat 0x01 - General Failure   
    
    perror("Neuspecny connect");
    shutdown(clientSock, SHUT_RDWR);
    close(clientSock);
    close(outSock);
    exit(1); // FIXME exit code    
  }
}

void reqBind(int clientSock,SocksRequest * req)
{
   
  sockaddr_in recipient = getRecipient(clientSock,req);   
  int outSock = socket(AF_INET, SOCK_STREAM, 0);  
  
  if(bind(outSock,(sockaddr*)&recipient,sizeof(recipient)) <0)
  {
    close(outSock);
    perror("Neuspesny bind.");
    exit(1);    
  }
  
  sockaddr_in sa;
  size_t sa_len = sizeof(sa);
  if(getsockname(outSock,(sockaddr *)&sa,(socklen_t *)&sa_len)<0)   
  {
      perror("getsockname");
      exit(1);
  }   
  
  //nastavenie replay adresy(BIND.ADDR)    
  uint8_t* addr = (uint8_t * )&sa.sin_addr.s_addr;  
  
  //nastavenie replay portu(BIND.PORT)
  uint16_t port = sa.sin_port;
  
  //bind success reply
  writeReplay(clientSock,0,addr,4,port);
  
  //pochopil som to tak ze povolime len 1 connection,
  
  if(listen(outSock,1)<0)
  {
    //TODO: poslat sparavu
    close(outSock);
    perror("Neuspesny listen.");
    exit(1);
  }
  
  if((outSock = accept(outSock,(sockaddr *)&sa,(socklen_t *)&sa_len))<0) //pozor, tu sa prepise povodny socket
  {    
    //TODO: poslat sparavu    
    close(outSock);
    perror("Neuspesny accept.");
    exit(1);
  }  
  
  if(getpeername(outSock,(sockaddr *)&sa,(socklen_t *)&sa_len)<0)   
  {
      perror("getpeername");
      exit(1);
  }   

  //nastavenie replay adresy(BIND.ADDR)    
  addr = (uint8_t * )&sa.sin_addr.s_addr;  
  
  //nastavenie replay portu(BIND.PORT)
  port = sa.sin_port;
  
  //accept success reply
  writeReplay(clientSock,0,addr,4,port);
  
  exchangeData(clientSock,outSock);
    
}

void reqUDP(int clientSock, SocksRequest * req)
{
  
  // gets ready to send and recieve UDP messages
  int udpSock = socket(AF_INET, SOCK_DGRAM, 0);

  string addr(ADDRESS);

  int tr=1;	
  /*if(setsockopt(udpSock,SOL_SOCKET,SO_REUSEADDR,&tr,sizeof(int))<0)
  {
      perror("Chyba pri setsockopt");
  }*/

  // recipient structure holds every necessary recipient info
  sockaddr_in recipient = getRecipient(clientSock,req); 

  // connect instead of complicated bind
  if(connect(udpSock, (struct sockaddr *)&recipient, sizeof(struct sockaddr)) == -1){
    perror("UDP Associate - connect to udp recipient:");
    writeReplay(clientSock, 5);
    exit(1); // FIXME error code
  }   

  // success reply
  uint8_t replyAddrLen = 0;
  if(recipient.sin_family == AF_INET)
    replyAddrLen = 4;
  else if(recipient.sin_family == AF_INET6)
    replyAddrLen = 16;

  writeReplay(clientSock, 0, (uint8_t *)&recipient.sin_addr.s_addr , replyAddrLen , recipient.sin_port); // udpklient ju nevie spracovat, tak ju neposielam, ale to je jedno teraz
  puts("REPLY SENT");

  exchangeData(clientSock, udpSock, true);
}

int main(int argc, char* argv[])
{
	//defaults
	bool debug = false;
	string addr = ADDRESS;
	int port = PORT;	
	
	//argument parsing
	for(int i=1;i<argc;i++)	  	
	{
	  string s = argv[i];
	  if(s=="stop")
	  {	    
	    stop(extractFname(argv[0]));
	    return 0;
	  }
	  
	  if(s=="--debug")
	  {
	    debug = true;
	  }
	  
	  if(s=="--address")
	  {
	    if(++i>=argc)
	    {
	      printf("Bad parameters.");
	      return 2;
	    }
	    addr = argv[i];
	  }
	  
	  if(s=="--port")
	  {
	    if(++i>=argc)
	    {
	      printf("Bad parameters.");
	      return 2;
	    }
	    stringstream ss(argv[i]);
	    ss >> port;
	  }
	}
	
	//daemonisation
	if(!debug)
	{
	  if(daemon(1,0)<0)
	  {
	    perror("Chyba pri daemon.");
	  }
	}
	
	//creating socket
	int sck, sock;	
 	sockaddr_in address;
 	sockaddr address2;			
	socklen_t addr2_size = sizeof(address2);		
	sck = socket(PF_INET,SOCK_STREAM,0);	
	if(sck<0)
	{
	  perror("Nepodarilo sa vytvorit socket!");
	  return 1;
	}	
	
	//Prevent "Address already in use" message
	int tr=1;	
	if(setsockopt(sck,SOL_SOCKET,SO_REUSEADDR,&tr,sizeof(int))<0 && debug)
	{
	    perror("Chyba pri setsockopt");
	}
		
	//Binding socket
	address.sin_family = AF_INET;
	address.sin_addr.s_addr = inet_addr(addr.c_str());
	address.sin_port = htons(port);	
	if(bind(sck,(sockaddr*)&address,sizeof(address))<0)
	{
	    close(sck);
	    perror("Neuspesny bind.");
	    return 1;	
	}
	
	//set listen mode
	if(listen(sck,10)<0)
	{
	  close(sck);
	  perror("Neuspesny listen.");
	  return 1;
	}

        if(debug){
          printf("Host: %s\nPort: %d\n", ADDRESS, PORT);
        }
	
	//Handle SIGCHLD signal
	signal(SIGCHLD,childdied);
	
	//Accept connection
	while((sock = accept(sck,&address2,&addr2_size))>=0)
	{
	  //create child
	  int pid = fork();
	  if (pid<0)
	  {
	    perror("Chyba pri fork");
	    break;
	  }
	  //parent stays in listen mode
	  if(pid > 0) 
	  {
	    close(sock);
	    continue;
	  }	  
	  //child processes the request
	  close(sck);	  
          
          IdMessage * imsg = readIDMessage(sock);                    
          checkProtocol(sock,imsg->version);          
          //TODO ak prisiel protokol styri mozno sa oplati poslat 0xFF ako metodu
          writeIDReply(sock,imsg);    
          
          SocksRequest * req =  readRequest(sock);
          
          switch(req->cmd)
          {
            case 1://connect
              reqConnect(sock,req);
              break;
            case 2://bind
              reqBind(sock,req);
              break;
            case 3://UDP Associate
              reqUDP(sock, req);
              break;
          }
          
	  break;
	}	
	
	if(sock!=-1)
	{	   
	  char buf[BUF_SZ];	  
	  shutdown(sock,SHUT_WR);	
	  while(read(sock,buf,sizeof(buf))>0); // naco je dobry takyto riadok?
	  close(sock);
	  if(debug)
      puts("Done");
	}
	else
	{
	    perror("Chyba pri accept");
	    return 1;	 
	}
	
	return 0;
}
